package com.novunda.tool.codec;

import java.util.regex.Matcher;
import java.util.regex.Pattern;

class SM4Utils {

    private static final String CHARSET_UTF8 = "UTF-8";

    private String secretKey = "";
    private String iv = "";
    private boolean hexString = false;

    public SM4Utils() {
    }

    private boolean checkSecretKeyAndIv() throws Exception {
        if (iv == null || secretKey == null) {
            throw new Exception("iv or secretKey is null");
        }
        return true;
    }

    private boolean checkSecretKey() throws Exception {
        if (secretKey == null) {
            throw new Exception("secretKey is null");
        }
        return true;
    }

    public String encryptData_ECB(String plainText) {
        try {

            checkSecretKey();

            SM4Context ctx = new SM4Context();
            ctx.isPadding = true;
            ctx.mode = SM4.SM4_ENCRYPT;

            byte[] keyBytes;
            if (hexString) {
                keyBytes = Util.hexStringToBytes(secretKey);
            } else {
                keyBytes = secretKey.getBytes();
            }

            SM4 sm4 = new SM4();
            sm4.sm4_setkey_enc(ctx, keyBytes);
            byte[] encrypted = sm4.sm4_crypt_ecb(ctx, plainText.getBytes(CHARSET_UTF8));
            String cipherText = Base58.encode(encrypted);
            if (cipherText != null && cipherText.trim().length() > 0) {
                Pattern p = Pattern.compile("\\s*|\t|\r|\n");
                Matcher m = p.matcher(cipherText);
                cipherText = m.replaceAll("");
            }
            return cipherText;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    public String decryptData_ECB(String cipherText) {
        try {

            checkSecretKey();

            SM4Context ctx = new SM4Context();
            ctx.isPadding = true;
            ctx.mode = SM4.SM4_DECRYPT;

            byte[] keyBytes;
            if (hexString) {
                keyBytes = Util.hexStringToBytes(secretKey);
            } else {
                keyBytes = secretKey.getBytes();
            }

            SM4 sm4 = new SM4();
            sm4.sm4_setkey_dec(ctx, keyBytes);
            byte[] decrypted = sm4.sm4_crypt_ecb(ctx, Base58.decode(cipherText));
            return new String(decrypted, CHARSET_UTF8);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    public String encryptData_CBC(String plainText) {
        try {

            checkSecretKeyAndIv();

            SM4Context ctx = new SM4Context();
            ctx.isPadding = true;
            ctx.mode = SM4.SM4_ENCRYPT;

            byte[] keyBytes;
            byte[] ivBytes;
            if (hexString) {
                keyBytes = Util.hexStringToBytes(secretKey);
                ivBytes = Util.hexStringToBytes(iv);
            } else {
                keyBytes = secretKey.getBytes();
                ivBytes = iv.getBytes();
            }

            SM4 sm4 = new SM4();
            sm4.sm4_setkey_enc(ctx, keyBytes);
            byte[] encrypted = sm4.sm4_crypt_cbc(ctx, ivBytes, plainText.getBytes(CHARSET_UTF8));
            String cipherText = Base58.encode(encrypted);
            if (cipherText != null && cipherText.trim().length() > 0) {
                Pattern p = Pattern.compile("\\s*|\t|\r|\n");
                Matcher m = p.matcher(cipherText);
                cipherText = m.replaceAll("");
            }
            return cipherText;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    public String decryptData_CBC(String cipherText) {
        try {

            checkSecretKeyAndIv();

            SM4Context ctx = new SM4Context();
            ctx.isPadding = true;
            ctx.mode = SM4.SM4_DECRYPT;

            byte[] keyBytes;
            byte[] ivBytes;
            if (hexString) {
                keyBytes = Util.hexStringToBytes(secretKey);
                ivBytes = Util.hexStringToBytes(iv);
            } else {
                keyBytes = secretKey.getBytes();
                ivBytes = iv.getBytes();
            }

            SM4 sm4 = new SM4();
            sm4.sm4_setkey_dec(ctx, keyBytes);
            byte[] decrypted = sm4.sm4_crypt_cbc(ctx, ivBytes, Base58.decode(cipherText));
            return new String(decrypted, CHARSET_UTF8);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    public void setSecretKey(String secretKey) {
        if (secretKey.length() != 16) {
            this.secretKey = Str2MD5.MD5(secretKey);
        } else {
            this.secretKey = secretKey;
        }
    }

    public void setIv(String iv) {
        this.iv = iv;
    }
}
